use num_traits::{Num, NumCast};

use crate::nn::init::uniform;
use crate::nn::module::Module;
use crate::Tensor;

pub struct Linear<T> {
    weight: Tensor<T>,
    bias: Option<Tensor<T>>,
}

impl<T> Linear<T>
    where
        T: Copy + Num + NumCast + 'static,
{
    pub fn new(in_features: usize, out_features: usize, bias: bool) -> Linear<T> {
        let k = 1. / in_features as f64;

        let weight = uniform(&[in_features, out_features], -k.sqrt(), k.sqrt());
        weight.requires_grad_(true);

        let bias = if bias {
            let bias = uniform(&[out_features], -k.sqrt(), k.sqrt());
            bias.requires_grad_(true);
            Some(bias)
        } else {
            None
        };

        Self {
            weight,
            bias,
        }
    }
}

impl<T> Module<T> for Linear<T>
    where
        T: Copy + Num + 'static,
{
    fn forward(&self, tensor: Tensor<T>) -> Tensor<T> {
        match &self.bias {
            Some(bias) => &tensor.mm(&self.weight) + bias,
            None => tensor.mm(&self.weight),
        }
    }

    fn parameters(&self) -> Vec<Tensor<T>> {
        let mut params = vec![self.weight.clone()];
        if let Some(bias) = &self.bias {
            params.push(bias.clone())
        }
        params
    }
}